#include "RigidBodyRestingContactHandler.h"

namespace coolpool
{

	RigidBodyRestingContactHandler::RigidBodyRestingContactHandler()
	{

	}

	RigidBodyRestingContactHandler::~RigidBodyRestingContactHandler()
	{

	}

	void RigidBodyRestingContactHandler::handleRestingContact(const std::vector<const Collision *> * resting_contacts)
	{
		C_.clear();
		NC_.clear();
		f_ = new Eigen::VectorXd(resting_contacts->size());
		setZero(f_);
		A_ = computeAMatrix(resting_contacts);
		b_ = computeBVector(resting_contacts);
		Eigen::VectorXd B_starting = Eigen::VectorXd(*b_);

		int d = findLessThanZero(b_);
		while (d >= 0)
		{
			driveToZero(d);
			d = findLessThanZero(b_);
		}

		RigidBody * rigid_body1;
		RigidBody * rigid_body2;
		math_tools::Vector3D normal;
		math_tools::Vector3D point;

		Eigen::VectorXd final_check = (*A_) * (*f_) + (B_starting);
		for (unsigned i = 0; i < resting_contacts->size(); ++i)
		{

			normal = (*resting_contacts)[i]->collision_normal;
			point = (*resting_contacts)[i]->collision_point;

			rigid_body1 = static_cast<RigidBody *> ((*resting_contacts)[i]->first);
			rigid_body2 = static_cast<RigidBody *> ((*resting_contacts)[i]->second);

			addFrictionToImpulse((*resting_contacts)[i], rigid_body1, rigid_body2, (*f_)[i]);
			rigid_body1->addForce(normal * (*f_)[i]);

			// TODO: Change the angular velocity correctly
			//	body1->addTorque((point - body1->getCenterOfMassLocation()) * (*f_)[i]);

			if (!rigid_body2->isStatic())
			{
				rigid_body2->addForce((-normal) * (*f_)[i]);
				//	body2->addTorque(-((point - body1->getCenterOfMassLocation()) * (*f_)[i]));
			}
		}

		delete A_;
		delete b_;
		delete f_;
	}

	double RigidBodyRestingContactHandler::computeAij(const Collision * ci, const Collision * cj)
	{
		if ((ci->first != cj->first) && (ci->second != cj->second) && (ci->first != cj->second) && (ci->second != cj->first))
		{
			return 0.0;
		}

		RigidBody * body1 = static_cast<RigidBody *> (ci->first);
		RigidBody * body2 = static_cast<RigidBody *> (ci->second);

		math_tools::Vector3D ni = ci->collision_normal;
		math_tools::Vector3D nj = cj->collision_normal;
		math_tools::Vector3D pi = ci->collision_point;
		math_tools::Vector3D pj = cj->collision_point;
		math_tools::Vector3D ra = body1->getPointRelativeLocation(pi);
		math_tools::Vector3D rb = body2->getPointRelativeLocation(pi);

		math_tools::Vector3D force_on_body1;
		math_tools::Vector3D torque_on_body1;

		if (cj->first == ci->first)
		{
			force_on_body1 = nj;
			torque_on_body1 = (pj - body1->getCenterOfMassLocation()) * nj;
		}
		else if (cj->second == ci->first)
		{
			force_on_body1 = -nj;
			torque_on_body1 = (pj - body1->getCenterOfMassLocation()) * nj;
		}

		math_tools::Vector3D body1_linear = force_on_body1 * (1 / body1->getMass());
		math_tools::Vector3D tmp;
		math_tools::Vector3D body1_angular = tmp * ra;

		assert(!body1->isStatic());

		if (body2->isStatic())
		{
			return ni.dot((body1_linear + body1_angular));
		}
		math_tools::Vector3D force_on_body2;
		math_tools::Vector3D torque_on_body2;

		if (cj->first == ci->second)
		{
			force_on_body2 = nj;
			torque_on_body2 = (pj - body2->getCenterOfMassLocation()) * nj;
		}
		else if (cj->second == ci->second)
		{
			force_on_body2 = -nj;
			torque_on_body2 = (pj - body2->getCenterOfMassLocation()) * nj;
		}

		math_tools::Vector3D body2_linear = force_on_body2 * (1 / body2->getMass());
		matrix_tools::transform(body2->getInertiaTensorInverse(), torque_on_body2, tmp);
		math_tools::Vector3D body2_angular = tmp * torque_on_body2;

		return ni.dot((body1_linear + body1_angular) - (body2_linear + body2_angular));
	}

	const Eigen::MatrixXd * RigidBodyRestingContactHandler::computeAMatrix(const vector<const Collision*> * resting_contacts)
	{
		int count = resting_contacts->size();
		Eigen::MatrixXd * A = new Eigen::MatrixXd(count, count);
		for (int i = 0; i < count; ++i)
		{
			for (int j = 0; j < count; ++j)
			{
				(*A)(i, j) = computeAij((*resting_contacts)[i], (*resting_contacts)[j]);
			}
		}

		return A;
	}

	double RigidBodyRestingContactHandler::computeBElementStatic(const Collision * collision)
	{
		if (collision->first->isStatic() || !collision->second->isStatic())
			throw UnexpectedInputException();

		RigidBody * body = static_cast<RigidBody *> (collision->first);
		math_tools::Vector3D collision_normal = collision->collision_normal;
		math_tools::Vector3D external_torque = body->getTorque();
		matrix_tools::transform(body->getInertiaTensorInverse(), body->getTorque(), external_torque);
		math_tools::Vector3D ext_part = external_torque + body->getForce() * (1.0 / body->getMass());
		return collision_normal.dot(body->pointVelocity(collision->collision_point) + ext_part);
	}

	double RigidBodyRestingContactHandler::computeBElementDynamic(const Collision * collision)
	{
		if (collision->first->isStatic() || collision->second->isStatic())
			throw UnexpectedInputException();

		RigidBody * body1 = static_cast<RigidBody *> (collision->first);
		RigidBody * body2 = static_cast<RigidBody *> (collision->second);
		math_tools::Vector3D collision_normal = collision->collision_normal;

		math_tools::Vector3D external_torque1;
		matrix_tools::transform(body1->getInertiaTensorInverse(), body1->getTorque(), external_torque1);
		math_tools::Vector3D external_part1 = body1->getForce() * (1.0 / body1->getMass()) + external_torque1;

		math_tools::Vector3D external_torque2;
		matrix_tools::transform(body2->getInertiaTensorInverse(), body2->getTorque(), external_torque2);
		math_tools::Vector3D external_part2 = body2->getForce() * (1.0 / body2->getMass()) + external_torque2;

		return collision_normal.dot((body1->pointVelocity(collision->collision_point) + external_part1) - (body2->pointVelocity(collision->collision_point) + external_part2));

		///////////////////////////////////////////////////////////////////////////////////////////////////
		//////// Implementation for using differenciations ////////////////////////////////////////////////
		///////////////////////////////////////////////////////////////////////////////////////////////////
		//		math_tools::Vector3D ra = body1->getPointRelativeLocation(collision->collision_point);
		//		math_tools::Vector3D rb = body2->getPointRelativeLocation(collision->collision_point);
		//
		//		math_tools::Vector3D ext_part_body1 = ...forces...
		//
		//		math_tools::Vector3D tmp;
		//		math_tools::Vector3D ang_vel_body1 = body1->getAngularVelocity();
		//		matrix_tools::transform(body1->getInertiaTensorInverse(), body1->getAngularMomentum() * body1->getAngularVelocity(), tmp);
		//		math_tools::Vector3D vel_part_body1 = (ang_vel_body1 * (ang_vel_body1 * ra)) + (tmp * ra);
		//
		//		math_tools::Vector3D ext_part_body2 = ...forces...;
		//
		//		math_tools::Vector3D ang_vel_body2 = body2->getAngularVelocity();
		//		matrix_tools::transform(body2->getInertiaTensorInverse(), body2->getAngularMomentum() * body2->getAngularVelocity(), tmp);
		//		math_tools::Vector3D vel_part_body2 = (ang_vel_body2 * (ang_vel_body2 * rb)) + (tmp * rb);
		//
		//		math_tools::Vector3D n_dot = computeNormalChange(collision);
		//		double k2 = 2 * n_dot.dot(body1->pointVelocity(collision->collision_point) - body2->pointVelocity(collision->collision_point));
		//		return k1 + k2;
		////////////////////////////////////////////////////////////////////////////////////////////////////
	}

	Eigen::VectorXd * RigidBodyRestingContactHandler::computeBVector(const vector<const Collision *> * resting_contacts)
	{
		Eigen::VectorXd * vector_b = new Eigen::VectorXd(resting_contacts->size());
		const Collision * current;
		for (unsigned i = 0; i < resting_contacts->size(); ++i)
		{
			current = (*resting_contacts)[i];
			(*vector_b)[i] = current->second->isStatic() ? computeBElementStatic(current) : computeBElementDynamic(current);
		}
		return vector_b;
	}

	math_tools::Vector3D RigidBodyRestingContactHandler::computeNormalChange(const Collision * col)
	{
		return col->second->isStatic() ? math_tools::Vector3D(0.0, 0.0, 0.0) : static_cast<RigidBody *> (col->second)->getAngularVelocity() * col->collision_normal;
	}

	void RigidBodyRestingContactHandler::driveToZero(int d)
	{
		while (true)
		{
			Eigen::VectorXd delta_f = fdirection(d);
			Eigen::VectorXd delta_a = (*A_) * delta_f;
			pair<double, int> p = maxstep(*f_, *b_, delta_f, delta_a, d);
			double s = p.first;
			int j = p.second;
			(*f_) = (*f_) + delta_f * s;
			(*b_) = (*b_) + delta_a * s;

			if (C_.find(j) != C_.end())
			{
				C_.erase(j);
				NC_.insert(j);
			}
			else if (NC_.find(j) != NC_.end())
			{
				NC_.erase(j);
				C_.erase(j);
			}
			else
			{
				C_.insert(j);
				return;
			}
		}
	}

	Eigen::VectorXd RigidBodyRestingContactHandler::fdirection(int d)
	{
		Eigen::VectorXd delta_f(b_->rows());
		setZero(&delta_f);
		delta_f[d] = 1.0;

		if (C_.size() == 0)
			return delta_f;

		Eigen::MatrixXd * A_cc = sliceMatrix(*A_, C_);
		Eigen::VectorXd * v1 = sliceVector(A_->block(0, d, A_->cols(), 1), C_);

		Eigen::VectorXd x = A_cc->fullPivHouseholderQr().solve(-(*v1));
		delete A_cc;
		delete v1;
		return transferTof(delta_f, x);
	}

	int RigidBodyRestingContactHandler::findLessThanZero(Eigen::VectorXd * a)
	{
		for (int i = 0; i < a->rows(); ++i)
		{
			if ((*a)[i] < -0.00001)
				return i;
		}
		return -1;
	}

	std::pair<double, int> RigidBodyRestingContactHandler::maxstep(const Eigen::VectorXd & f, const Eigen::VectorXd & a, const Eigen::VectorXd & delta_f, const Eigen::VectorXd & delta_a, int d)
	{
		double s = numeric_limits<double>::max();
		int j = -1;
		if (delta_a[d] > 0)
		{
			j = d;
			s = -a[d] / delta_a[d];
		}
		for (std::set<int>::iterator i = C_.begin(); i != C_.end(); ++i)
		{
			if (delta_f[*i] < 0)
			{
				double s_ = -f[*i] / delta_f[*i];
				if (s_ < s)
				{
					s = s_;
					j = *i;
				}
			}
		}
		for (std::set<int>::iterator i = NC_.begin(); i != NC_.end(); ++i)
		{
			if (delta_a[*i] < 0)
			{
				double s_ = -a[*i] / delta_a[*i];
				if (s_ < s)
				{
					s = s_;
					j = *i;
				}
			}
		}

		return pair<double, int> (s, j);
	}

	Eigen::VectorXd RigidBodyRestingContactHandler::transferTof(Eigen::VectorXd f, Eigen::VectorXd & vector)
	{
		int i = 0;
		for (std::set<int>::iterator row = C_.begin(); row != C_.end(); ++row)
		{
			f[*row] = vector[i++];
		}

		return f;
	}
}
